from const import *
from typing import List
from random import uniform
from math import sqrt
from utils import *


class Ball:
    '''
    球体类，模拟球体在桌面上的运动
    '''


    # 速度单位：像素/秒
    def __init__(self, color = (0, 0, 0), x = 100, y = 100, vx = 0, vy = 0, radius = 15, m = 1):
        self.m = m  # 质量
        self.color = color  # 颜色
        self.radius = radius  # 半径

        self.x, self.y = x, y  # 坐标
        self.last_x, self.last_y = x, y  # 上一帧的坐标
        self.vx, self.vy = vx, vy  # 速度
        self.move_x, self.move_y = True, True  # 是否在运动
        self.ratio = abs(self.vx / self.vy) if abs(self.vy) > 0 else 1  # x轴上摩擦力的分力比值

        if self.x - self.radius < TABLE_LEFT or self.x + self.radius > TABLE_RIGHT or self.y - self.radius < TABLE_UP \
                or self.y + self.radius > TABLE_BOTTOM:
            raise ('generate the ball default.')


    def position(self, t, ):
        '''
        实时定位球体的位置坐标
        :param t: 运动时间
        :return: None
        '''
        if self.move_x:
            self.last_x = self.x
            self.x = self.x + self.vx * t + 0.5 * self.miu_x * g * t * t
        if self.move_y:
            self.last_y = self.y
            self.y = self.y + self.vy * t + 0.5 * self.miu_y * g * t * t


    def velocity(self, t, miu = miu):
        '''
        实时计算球体当前速度
        :param t: 运动时间
        :param miu: 摩擦系数
        :return: None
        '''
        # miu = round(uniform(0.95, 1.05), 2) * miu
        self.miu_x = miu if self.vx > 0 else -miu
        self.miu_y = miu if self.vy > 0 else -miu

        # if self.vx < 10 or self.vy < 10:
        #     self.miu_x *= 3
        #     self.miu_y *= 3

        if self.vx != 0 and self.move_x:
            temp_vx = self.vx + self.ratio * self.miu_x * g * t
            if abs(self.vx) > abs(temp_vx):
                self.vx = temp_vx
            else:
                self.vx = 0
                self.move_x = False
        else:
            self.move_x = False

        if self.vy != 0 and self.move_y:
            temp_vy = self.vy + self.miu_y * g * t
            if abs(self.vy) > abs(temp_vy):
                self.vy = temp_vy
            else:
                self.vy = 0
                self.move_y = False
        else:
            self.move_y = False


    def bump(self):
        '''
        球体撞到桌沿，损失一定的机械能
        :return: None
        '''
        left, right, up, bottom = TABLE_LEFT, TABLE_RIGHT, TABLE_UP, TABLE_BOTTOM
        # 上
        if self.y - self.radius < up:
            # print('up:',self.vx,self.vy)
            self.vy *= -(1-LOSS_RATIO1)
            self.vx *= 1-LOSS_RATIO1
            self.y = up+self.radius+1
        # 下
        elif self.y + self.radius > bottom:
            # print('bottom:',self.vx,self.vy)
            self.vy *= -(1-LOSS_RATIO1)
            self.vx *= 1-LOSS_RATIO1
            self.y = bottom-self.radius-1
        # 左
        elif self.x - self.radius < left:
            # print('left:',self.vx,self.vy)
            self.vx *= -(1-LOSS_RATIO1)
            self.vy *= 1-LOSS_RATIO1
            self.x = left+self.radius+1
        # 右
        elif self.x + self.radius > right:
            # print('right:',self.vx,self.vy)
            self.vx *= -(1-LOSS_RATIO1)
            self.vy *= 1-LOSS_RATIO1
            self.x = right-self.radius-1


    def collide(self, all_balls: list, collide_balls: dict):
        '''
        两球相碰
        :param all_balls: 所有球体
        :param collide_balls: 碰撞过的球体对，防止多次判断碰撞
        :return: collide_balls
        '''
        for i in all_balls:
            last_dist = sqrt((self.last_x - i.last_x) ** 2 + (self.last_y - i.last_y) ** 2)
            # dist=sqrt((self.x - i.x) ** 2 + (self.y - i.y) ** 2)
            dx = i.x - self.x
            dy = i.y - self.y
            dist = sqrt(dx ** 2 + dy ** 2)
            if i != self and last_dist > self.radius + i.radius and dist <= self.radius + i.radius:
                for k, v in collide_balls.items():
                    if self in v:
                        return collide_balls
                else:
                    if collide_balls.get(self):
                        collide_balls[self].append(i)
                    else:
                        collide_balls[self] = [i]

                    # 碰撞损失机械能
                    self.vx *= 1-LOSS_RATIO2
                    self.vy *= 1-LOSS_RATIO2

                    nx = dx / dist
                    ny = dy / dist
                    tx = -ny
                    ty = nx
                    dpTan1 = self.vx * tx + self.vy * ty
                    dpTan2 = i.vx * tx + i.vy * ty
                    dpNorm1 = self.vx * nx + self.vy * ny
                    dpNorm2 = i.vx * nx + i.vy * ny
                    m1 = (dpNorm1 * (self.m - i.m) + 2 * i.m * dpNorm2) / (self.m + i.m)
                    m2 = (dpNorm2 * (i.m - self.m) + 2 * self.m * dpNorm1) / (self.m + i.m)
                    self.vx = tx * dpTan1 + nx * m1
                    self.vy = ty * dpTan1 + ny * m1
                    i.vx = tx * dpTan2 + nx * m2
                    i.vy = ty * dpTan2 + ny * m2

                    i.move_x, self.move_x, i.move_y, self.move_y = True, True, True, True

        return collide_balls


if __name__ == '__main__':
    ball = Ball()
